# 实现常用函数的计算
import time
import numpy as np
from dolfinx import geometry, mesh
import numpy as np
from scipy.sparse import csc_matrix
from scipy.sparse.linalg import spsolve
from scipy.special import legendre


def sparse_solve(A, b, print_=True):
    """
    将输入的矩阵A和向量b转换为稀疏格式并求解线性方程组Ax=b

    参数:
    A (np.ndarray): 系数矩阵
    b (np.ndarray): 右侧向量
    print_: 打印稀疏线性系统求解时间
    返回:
    np.ndarray: 方程组的解
    """
    # 转换为稀疏矩阵格式（CSC格式适合求解）
    start_time = time.time()
    A_sparse = csc_matrix(A)

    # 调用稀疏求解器
    res = spsolve(A_sparse, b)
    end_time = time.time()
    if print_:
        print(
            f"Time taken to solve the system: {end_time - start_time:.4f} seconds")
    return res


class GaussIntegrate1D:
    def __init__(self, order=9):
        """初始化高斯积分器"""
        assert order >= 3, "积分阶数过低!"
        self.order = order
        self.points, self.weights = self._get_quadrature(order)

    def _get_quadrature(self, order):
        """计算标准区间 [-1, 1] 上的高斯积分点和权重"""
        n = (order + 1) // 2
        poly = np.polynomial.Legendre([0]*(n) + [1])
        roots = poly.roots()  # 积分点
        dpoly = poly.deriv()
        weights = 2.0 / ((1.0 - roots**2) * (dpoly(roots))**2)
        return roots, weights

    def get_nodes(self, a=-1, b=1):
        """
        返回指定区间 [a, b] 上的积分节点坐标

        参数:
            a (float): 积分区间起点，默认为 -1
            b (float): 积分区间终点，默认为 1

        返回:
            np.ndarray: 积分节点坐标数组
        """
        # 将标准区间 [-1, 1] 上的节点映射到 [a, b]
        return 0.5 * (b - a) * self.points + 0.5 * (a + b)

    def integrate(self, a, b, values):
        """
        使用预计算的节点值进行积分

        参数:
            a (float): 积分区间起点
            b (float): 积分区间终点
            values (np.ndarray): 对应 get_nodes(a, b) 返回节点的函数值

        返回:
            float: 积分结果
        """
        if len(values) != len(self.points):
            raise ValueError(f"values 长度必须为 {len(self.points)}，与积分阶数匹配")

        # 缩放权重以适应区间 [a, b]
        scaled_weights = 0.5 * (b - a) * self.weights
        return np.sum(scaled_weights * values)

# 存储高斯积分点


class GaussIntegrate:
    def __init__(self, order=3):
        # 初始化函数，设置默认阶数为3
        self.order = order
        # 调用_get_quadrature函数，获取积分点和权重
        self.points, self.weights = self._get_quadrature(order)

    def _get_quadrature(self, order):
        """获取三角形高斯积分点和权重
            order 为高斯积分点数
        """
        if order == 3:  # 3点积分，精确积分二次多项式
            points = np.array([
                [0, 0.5],
                [0.5, 0],
                [0.5, 0.5]
            ])
            weights = np.array([1/6, 1/6, 1/6])

        elif order == 4:  # 4点积分，精确积分三次多项式
            sqrt3 = np.sqrt(3)
            points = np.array([
                [(1 + sqrt3/3)/2, (1 - sqrt3/3) * (1 + sqrt3/3)/4],
                [(1 + sqrt3/3)/2, (1 - sqrt3/3) ** 2 / 4],
                [(1 - sqrt3/3)/2, (1 + sqrt3/3) ** 2 / 4],
                [(1 - sqrt3/3)/2, (1 + sqrt3/3) * (1 - sqrt3/3)/4]
            ])
            weights = np.array([
                (1 - sqrt3/3)/8,
                (1 - sqrt3/3)/8,
                (1 + sqrt3/3)/8,
                (1 + sqrt3/3)/8
            ])

        elif order == 7:  # 7点积分，精确积分五次多项式
            sqrt15 = np.sqrt(15)
            points = np.array([
                [1/3, 1/3],
                [(6 - sqrt15)/21, (6 - sqrt15)/21],
                [(6 - sqrt15)/21, (9 + 2*sqrt15)/21],
                [(9 + 2*sqrt15)/21, (6 - sqrt15)/21],
                [(6 + sqrt15)/21, (6 + sqrt15)/21],
                [(6 + sqrt15)/21, (9 - 2*sqrt15)/21],
                [(9 - 2*sqrt15)/21, (6 + sqrt15)/21]
            ])
            weights = np.array([
                0.1125,
                (155 - sqrt15)/2400,
                (155 - sqrt15)/2400,
                (155 - sqrt15)/2400,
                (155 + sqrt15)/2400,
                (155 + sqrt15)/2400,
                (155 + sqrt15)/2400
            ])

        elif order == 9:  # 9点积分，精确积分五次多项式（更高精度）
            sqrt3_5 = np.sqrt(3/5)
            points = np.array([
                [0.5, 0.25],
                [(1 + sqrt3_5)/2, (1 - sqrt3_5) * (1 + sqrt3_5)/4],
                [(1 + sqrt3_5)/2, (1 - sqrt3_5) ** 2 / 4],
                [(1 - sqrt3_5)/2, (1 + sqrt3_5) ** 2 / 4],
                [(1 - sqrt3_5)/2, (1 + sqrt3_5) * (1 - sqrt3_5)/4],
                [0.5, (1 - 0) * (1 + sqrt3_5)/4],
                [0.5, (1 - 0) * (1 - sqrt3_5)/4],
                [(1 + sqrt3_5)/2, (1 - sqrt3_5) * (1 + 0)/4],
                [(1 - sqrt3_5)/2, (1 + sqrt3_5) * (1 + 0)/4]
            ])
            weights = np.array([
                64/81 * (1-0)/8,
                100/324 * (1 - sqrt3_5)/8,
                100/324 * (1 - sqrt3_5)/8,
                100/324 * (1 + sqrt3_5)/8,
                100/324 * (1 + sqrt3_5)/8,
                40/81 * (1-0)/8,
                40/81 * (1-0)/8,
                40/81 * (1 - sqrt3_5)/8,
                40/81 * (1 + sqrt3_5)/8
            ])

        else:
            raise ValueError(f"Unsupported quadrature order: {order}. "
                             f"Supported orders: 3, 4, 7, 9")

        return points, weights

    def to_physical_gauss_points(self, p0, p1, p2):
        """
        将标准三角形的高斯点坐标和权重转换为物理三角形坐标和权重

        参数:
        p0, p1, p2: 物理三角形的三个顶点，每个顶点为(x, y)坐标对

        返回:
        physical_points: 物理三角形上的高斯点坐标
        physical_weights: 物理三角形上的高斯积分权重
        """
        # 提取顶点坐标
        x1, y1 = p0
        x2, y2 = p1
        x3, y3 = p2

        # 计算雅可比行列式（平行四边形面积）
        jacobi = abs((x2 - x1) * (y3 - y1) - (x3 - x1) * (y2 - y1))

        # 转换高斯点坐标到物理坐标
        # 使用公式：x = x1 + (x2-x1)*L1 + (x3-x1)*L2
        #          y = y1 + (y2-y1)*L1 + (y3-y1)*L2
        # 其中L1和L2是参考三角形上的坐标（self.points）
        L1 = self.points[:, 0]
        L2 = self.points[:, 1]
        physical_x = x1 + (x2 - x1) * L1 + (x3 - x1) * L2
        physical_y = y1 + (y2 - y1) * L1 + (y3 - y1) * L2
        physical_points = np.column_stack((physical_x, physical_y))

        # 计算物理权重 = 参考权重 * 雅可比行列式
        # physical_weights = self.weights * jacobi

        return physical_points

    def integrate(self, func, vertices):
        """
        在三角形上计算函数积分

        参数:
            func: 被积函数 f(x, y)
            vertices: 三角形顶点数组 [[x0,y0], [x1,y1], [x2,y2]]

        返回:
            积分值
        """
        p0, p1, p2 = vertices

        # 计算雅可比矩阵的行列式 (2倍三角形面积)
        J = np.array([
            [p1[0] - p0[0], p2[0] - p0[0]],
            [p1[1] - p0[1], p2[1] - p0[1]]
        ])
        detJ = abs(np.linalg.det(J))

        integral = 0.0
        for point, weight in zip(self.points, self.weights):
            xi, eta = point
            # 计算物理坐标
            x = p0[0] + xi*(p1[0]-p0[0]) + eta*(p2[0]-p0[0])
            y = p0[1] + xi*(p1[1]-p0[1]) + eta*(p2[1]-p0[1])

            # 累加积分
            integral += weight * func(x, y)

        return detJ * integral

    def get_vector_basis_values(self, num_basis_func, element):
        """返回矢量单元的基函数值

        Args:
            num_basis_func (_type_): 基函数的个数
            element (_type_): 有限元单元

        Returns:
            _type_: x分量和y分量的值
        """
        v = element.tabulate(0, self.points)[
            0]  # 尺寸为num_gs_points, 2, num_basis_func
        v_x = np.zeros((len(self.weights), num_basis_func))
        v_y = np.zeros((len(self.weights), num_basis_func))
        for i in range(len(self.weights)):
            v_x[i] = v[i, 0, ::2]
            v_y[i] = v[i, 1, 1::2]
        return v_x, v_y


class BasisDerivativeTransformer:
    """
    将基函数由标准单元变换到参考单元
    """

    def __init__(self, nodes: np.ndarray):
        """传入三角形顶点坐标数组 3行两列

        Args:
            nodes (np.ndarray): 三角形顶点坐标
        """
        p0, p1, p2 = nodes
        self.nodes = nodes
        self.p0 = p0
        self.p1 = p1
        self.p2 = p2
        self.J = self.jacobi_matrix()
        self.J11 = self.J[0][0]
        self.J12 = self.J[0][1]
        self.J21 = self.J[1][0]
        self.J22 = self.J[1][1]
        self.detJ = abs(np.linalg.det(self.J))

    def jacobi_matrix(self):
        """
        计算三角形的雅可比矩阵

        返回:
        2x2的雅可比矩阵
        """
        x0, y0 = self.p0
        x1, y1 = self.p1
        x2, y2 = self.p2

        # 计算雅可比矩阵的元素
        dx_xi = x1 - x0
        dx_eta = x2 - x0
        dy_xi = y1 - y0
        dy_eta = y2 - y0

        # 构建雅可比矩阵
        J = np.array([
            [dx_xi, dx_eta],
            [dy_xi, dy_eta]
        ])

        return J  # 请参考何晓明有限元程序

    def local_to_ref(self, x: np.ndarray, y: np.ndarray):
        """将local单元坐标变换到ref单元

        Args:
            x (np.ndarray): x坐标序列 1D vec
            y (np.ndarray): y坐标序列 1D vec
        """
        assert x.shape == y.shape, "x和y需要形状相同"
        coords = np.vstack([x, y])
        coords[0] -= self.p0[0]
        coords[1] -= self.p0[1]
        return np.linalg.inv(self.J) @ coords         

    def local(self, ref_value: np.ndarray):
        """计算 local单元基函数值

        Args:
            ref_value (np.ndarray): 标准单元上的值
        """
        return ref_value

    def local_dx1(self, ref_dx1: np.ndarray, ref_dy1: np.ndarray):
        """计算local dx的值

        Args:
            ref_dx1 (np.ndarray): 参考单元 对x求一阶导的结果
            ref_dy1 (np.ndarray): 参考单元 对y求一阶导的结果
        """
        r = (self.J22 * ref_dx1 - self.J21 * ref_dy1) / self.detJ
        return r

    def local_dy1(self, ref_dx1: np.ndarray, ref_dy1: np.ndarray):
        """计算local dy的值

        Args:
            ref_dx1 (np.ndarray): 参考单元 对x求一阶导的结果
            ref_dy1 (np.ndarray): 参考单元 对y求一阶导的结果
        """
        r = (-self.J12 * ref_dx1 + self.J11 * ref_dy1) / self.detJ
        return r

    def local_dx2(self, ref_dx2: np.ndarray, ref_dy2: np.ndarray, ref_dx1dy1: np.ndarray):
        """计算local x方向二阶导数

        Args:
            ref_dx2 (np.ndarray): 参考单元 对x求2阶导的结果
            ref_dy2 (np.ndarray): 参考单元 对y求2阶导的结果
            ref_dx1dy1 (np.ndarray): 参考单元 对x, y各求一阶导的结果
        """
        A = self.J22**2 * ref_dx2
        B = self.J21**2 * ref_dy2
        C = -2*self.J21 * self.J22 * ref_dx1dy1
        return (A + B + C) / self.detJ**2

    def local_dy2(self, ref_dx2: np.ndarray, ref_dy2: np.ndarray, ref_dx1dy1: np.ndarray):
        """计算local y方向二阶导数

        Args:
            ref_dx2 (np.ndarray): 参考单元 对x求2阶导的结果
            ref_dy2 (np.ndarray): 参考单元 对y求2阶导的结果
            ref_dx1dy1 (np.ndarray): 参考单元 对x, y各求一阶导的结果
        """
        A = self.J12**2 * ref_dx2
        B = self.J11**2 * ref_dy2
        C = -2*self.J11 * self.J12 * ref_dx1dy1
        return (A + B + C) / self.detJ**2

    def local_dx1dy1(self, ref_dx2: np.ndarray, ref_dy2: np.ndarray, ref_dx1dy1: np.ndarray):
        """计算local x和y方向各自求一阶导

        Args:
            ref_dx2 (np.ndarray): 参考单元 对x求2阶导的结果
            ref_dy2 (np.ndarray): 参考单元 对y求2阶导的结果
            ref_dx1dy1 (np.ndarray): 参考单元 对x, y各求一阶导的结果
        """
        A = -self.J22*self.J12 * ref_dx2
        B = -self.J11*self.J21 * ref_dy2
        C = (self.J11*self.J22 + self.J21*self.J12) * ref_dx1dy1
        return (A + B + C) / self.detJ**2



def reference_to_physical(element, p0, p1, p2, X):
    """
    传入有限元 三角形的三个顶点 标准单元坐标
    Map coordinates `X` in reference element to triangle defined by `p0`, `p1` and `p2`
    """
    # 计算参考单元基函数在 X 处的值（0 阶导数，即基函数本身）
    basis_values = element.tabulate(0, X)
    # 基函数加权物理顶点，得到物理坐标
    x = basis_values[0] @ np.vstack([p0, p1, p2])
    return x


def search_element(coords: np.ndarray, domain: mesh.Mesh, bb_tree: geometry.BoundingBoxTree):
    """传入坐标列表 返回对应的单元索引列表

    Args:
        coords (np.ndarray): 坐标列表 n*2
        domain (mesh.Mesh): 需要查询的区域
        bb_tree (geometry.BoundingBoxTree): 边界包围盒树

    Returns:
        _type_: 单元索引列表 list
    """
    res = []
    if coords.ndim == 1:
        coords = np.append(coords, 0)
        candidates = geometry.compute_collisions_points(bb_tree, coords)
        target_cell = geometry.compute_colliding_cells(
            domain, candidates, coords)
        return target_cell
    else:
        for coord in coords:
            # 二维需要补0
            coord = np.append(coord, 0)
            # 查询可能包含点的单元
            candidates = geometry.compute_collisions_points(bb_tree, coord)
            target_cell = geometry.compute_colliding_cells(
                domain, candidates, coord)
            # target_cell_id = target_cell.array[0]
            res.append(target_cell)
        return res


def process_case1_search_cell(res):
    count = {}
    for i in res:
        for v in i.array:
            if v not in count:
                count[v] = 1
            else:
                count[v] += 1
    # 返回出现次数最多的
    max_key = max(count, key=count.get)
    return max_key


def assign_to_corner(A, B, corner='lower_left'):
    """
    将矩阵B赋值到矩阵A的指定角落区域

    参数:
    A (np.ndarray): 目标矩阵
    B (np.ndarray): 源矩阵
    corner (str): 角落位置，可选值为 'upper_left', 'lower_left', 'upper_right', 'lower_right'

    返回:
    np.ndarray: 修改后的矩阵A
    """
    # 计算源矩阵B的行数和列数
    rows, cols = B.shape

    # 验证目标区域是否合法
    if rows > A.shape[0] or cols > A.shape[1]:
        raise ValueError(f"矩阵B的形状({rows},{cols})超出了矩阵A的范围({A.shape})")

    # 根据指定角落确定切片索引
    if corner == 'upper_left':
        A[:rows, :cols] = B
    elif corner == 'lower_left':
        start_row = A.shape[0] - rows
        A[start_row:, :cols] = B
    elif corner == 'upper_right':
        start_col = A.shape[1] - cols
        A[:rows, start_col:] = B
    elif corner == 'lower_right':
        start_row = A.shape[0] - rows
        start_col = A.shape[1] - cols
        A[start_row:, start_col:] = B
    else:
        raise ValueError(
            f"不支持的角落位置: {corner}，可选值为 'upper_left', 'lower_left', 'upper_right', 'lower_right'")

    return A


def search_boundary_edge_cell(msh: mesh.Mesh, counterclock=False) -> dict:
    """
    识别并返回有限元网格中所有边界边及其对应的单元索引和坐标

    参数:
        msh (mesh.Mesh): DOLFINx 网格对象
        counterclock (bool): 是否按逆时针排序边界边，默认为 False

    返回:
        dict: 字典，键为边界边的全局索引，值为包含单元索引和边坐标的字典
    """
    # 获取网格拓扑维度信息
    tdim = msh.topology.dim  # 网格维度 (2D=2)
    fdim = tdim - 1          # 边界维度 (2D=1)
    msh.topology.create_connectivity(fdim, tdim)
    msh.topology.create_connectivity(fdim, 0)  # 创建边到顶点的连接关系

    # 获取边到单元的连接关系
    facet_to_cells = msh.topology.connectivity(fdim, tdim)
    facet_to_vertices = msh.topology.connectivity(fdim, 0)  # 边到顶点的连接关系

    # 获取网格几何信息（顶点坐标）
    x = msh.geometry.x  # 所有顶点的坐标

    # 构建结果字典
    boundary_edges = {}
    for facet_index in range(facet_to_cells.num_nodes):
        cells = facet_to_cells.links(facet_index)
        if cells.size == 1:  # 边界边
            # 获取边的顶点索引
            vertices = facet_to_vertices.links(facet_index)
            # 获取顶点坐标
            coords = x[vertices][:, :tdim]  # 提取前tdim个坐标分量（通常是2D）

            boundary_edges[facet_index] = {
                'cell_index': cells[0],  # 连接的单元索引
                'coordinates': coords    # 边的顶点坐标（形状：(2, tdim)）
            }

    # 如果需要逆时针排序，进行后处理
    if counterclock and boundary_edges:
        # 计算所有边界边的质心作为参考点
        all_coords = np.vstack([edge['coordinates']
                               for edge in boundary_edges.values()])
        centroid = np.mean(all_coords, axis=0)

        # 对每条边界边，检查其方向是否为逆时针，否则反转
        for edge_data in boundary_edges.values():
            p1, p2 = edge_data['coordinates']
            # 计算向量从质心到p1和p2
            v1 = p1 - centroid
            v2 = p2 - centroid
            # 计算叉积判断方向（z分量）
            cross_product = v1[0] * v2[1] - v1[1] * v2[0]
            # 如果叉积为负，说明是顺时针方向，需要反转
            if cross_product < 0:
                edge_data['coordinates'] = np.array([p2, p1])

    return boundary_edges


def is_counter_clock(triangle: np.ndarray) -> bool:
    """
    判断三角形是否按逆时针方向排列
    
    参数:
        triangle: 3x2的numpy数组，每行表示一个顶点坐标
        
    返回:
        bool: 如果是逆时针方向返回True，否则返回False
    """
    # 验证输入格式
    assert triangle.shape == (3, 2), "输入必须是3x2的数组（3个顶点，每个顶点2个坐标）"
    
    # 提取三点坐标
    A = triangle[0] 
    B = triangle[1]
    C = triangle[2]
    
    # 计算向量AB和AC
    AB = B - A
    AC = C - A
    
    # 计算叉积（z分量，判断方向）
    cross_product = AB[0] * AC[1] - AB[1] * AC[0]
    return cross_product > 0  # 如果叉积大于0，则为逆时针方向



